from __future__ import division, print_function
import random

import numpy as np
import torch
import torch.nn as nn
from torch.distributions.uniform import Uniform


def sparse_init_weight(model):
    for m in model.modules():
        if isinstance(m, nn.Conv3d):
            torch.nn.init.sparse_(m.weight, sparsity=0.1)
        elif isinstance(m, nn.BatchNorm3d):
            m.weight.data.fill_(1)
            m.bias.data.zero_()
    return model


class ConvBlock(nn.Module):
    """two convolution layers with batch norm and leaky relu"""

    def __init__(self, in_channels, out_channels, dropout_p):
        super(ConvBlock, self).__init__()
        self.conv_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(),
            nn.Dropout(dropout_p),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(),
        )

    def forward(self, x):
        return self.conv_conv(x)


class DownBlock(nn.Module):
    """Downsampling followed by ConvBlock"""

    def __init__(self, in_channels, out_channels, dropout_p):
        super(DownBlock, self).__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2), ConvBlock(in_channels, out_channels, dropout_p)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class UpBlock(nn.Module):
    """Upssampling followed by ConvBlock"""

    def __init__(
        self, in_channels1, in_channels2, out_channels, dropout_p, bilinear=True
    ):
        super(UpBlock, self).__init__()
        self.bilinear = bilinear
        if bilinear:
            self.conv1x1 = nn.Conv2d(in_channels1, in_channels2, kernel_size=1)
            self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
        else:
            self.up = nn.ConvTranspose2d(
                in_channels1, in_channels2, kernel_size=2, stride=2
            )
        self.conv = ConvBlock(in_channels2 * 2, out_channels, dropout_p)

    def forward(self, x1, x2):
        if self.bilinear:
            x1 = self.conv1x1(x1)
        x1 = self.up(x1)
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class Encoder(nn.Module):
    def __init__(self, params):
        super(Encoder, self).__init__()
        self.params = params
        self.in_chns = self.params["in_chns"]
        self.ft_chns = self.params["feature_chns"]
        self.n_class = self.params["class_num"]
        self.bilinear = self.params["bilinear"]
        self.dropout = self.params["dropout"]
        assert len(self.ft_chns) == 5
        self.in_conv = ConvBlock(self.in_chns, self.ft_chns[0], self.dropout[0])
        self.down1 = DownBlock(self.ft_chns[0], self.ft_chns[1], self.dropout[1])
        self.down2 = DownBlock(self.ft_chns[1], self.ft_chns[2], self.dropout[2])
        self.down3 = DownBlock(self.ft_chns[2], self.ft_chns[3], self.dropout[3])
        self.down4 = DownBlock(self.ft_chns[3], self.ft_chns[4], self.dropout[4])

    def forward(self, x):
        x0 = self.in_conv(x)
        x1 = self.down1(x0)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        return [x0, x1, x2, x3, x4]


class Decoder(nn.Module):
    def __init__(self, params):
        super(Decoder, self).__init__()
        self.params = params
        self.in_chns = self.params["in_chns"]
        self.ft_chns = self.params["feature_chns"]
        self.n_class = self.params["class_num"]
        self.bilinear = self.params["bilinear"]
        assert len(self.ft_chns) == 5

        self.up1 = UpBlock(
            self.ft_chns[4], self.ft_chns[3], self.ft_chns[3], dropout_p=0.0
        )
        self.up2 = UpBlock(
            self.ft_chns[3], self.ft_chns[2], self.ft_chns[2], dropout_p=0.0
        )
        self.up3 = UpBlock(
            self.ft_chns[2], self.ft_chns[1], self.ft_chns[1], dropout_p=0.0
        )
        self.up4 = UpBlock(
            self.ft_chns[1], self.ft_chns[0], self.ft_chns[0], dropout_p=0.0
        )

        self.out_conv = nn.Conv2d(
            self.ft_chns[0], self.n_class, kernel_size=3, padding=1
        )

    def forward(self, feature):
        x0 = feature[0]
        x1 = feature[1]
        x2 = feature[2]
        x3 = feature[3]
        x4 = feature[4]

        x = self.up1(x4, x3)
        x = self.up2(x, x2)
        x = self.up3(x, x1)
        x = self.up4(x, x0)
        output = self.out_conv(x)
        return output


class PertDropout(nn.Module):
    def __init__(self, p=0.5):
        super(PertDropout, self).__init__()
        self.p = p
        self.dropouts = [
            nn.Dropout2d(p * 0.5).cuda(),  # Weak
            nn.Dropout2d(p * 1.5).cuda(),  # Strong
        ]
        self.len = len(self.dropouts)

    def __len__(self):
        return self.len

    def forward(self, x):
        rst = []
        for pert_dropout in self.dropouts:
            single_type = []
            for i, feat in enumerate(x):
                perted = pert_dropout(feat)
                single_type.append(perted)
            rst.append(single_type)
        return rst


class UNet(nn.Module):
    def __init__(self, in_chns, class_num):
        super(UNet, self).__init__()

        params = {
            "in_chns": in_chns,
            "feature_chns": [16, 32, 64, 128, 256],
            "dropout": [0.05, 0.1, 0.2, 0.3, 0.5],
            "class_num": class_num,
            "bilinear": False,
            "acti_func": "relu",
        }

        self.encoder = Encoder(params)
        self.decoder = Decoder(params)
        self.ema_decoder = None

        sparse_init_weight(self.encoder)
        sparse_init_weight(self.decoder)
        if self.ema_decoder is not None:
            sparse_init_weight(self.ema_decoder)

        self.pert = PertDropout(0.5).cuda()

    def forward(self, x, need_fp=False, need_ema=False, both=False, drop_rate=0.5):
        feature = self.encoder(x)

        if need_fp:
            features_x = []
            features_u = []
            for feats in feature:
                fx, fu = feats.chunk(2)
                features_x.append(fx)
                features_u.append(fu)

            perted_fus = self.pert(features_u)
            all_zip = zip(features_x, features_u, *perted_fus)
            outs = self.decoder([torch.cat(feats_all) for feats_all in all_zip])
            return outs.chunk(2 + len(self.pert))

        if need_ema:
            pert = (
                nn.FeatureAlphaDropout(0.5)
                if random.random() < 0.5
                else nn.AlphaDropout(0.5)
            )
            return self.decoder(feature), self.ema_decoder(
                [pert(feat) for feat in feature]
            )

        output = self.decoder(feature)
        return output


if __name__ == "__main__":
    model = UNet(1, 4).cuda()
    x = torch.randn((12, 1, 64, 64)).cuda()
    y = model(x, True)
    for i in y:
        print(i.shape)